#!/usr/bin/python
'''
consistently flips orientation of contigs in edge file.
'''
import sys
import random

# parameters.
node_file = sys.argv[1]
agp_file = sys.argv[2]
edge_file1 = sys.argv[3]
edge_file2 = sys.argv[4]
out1 = sys.argv[5]
out2 = sys.argv[6]
out3 = sys.argv[7]

# read in the node file.
flip_status = {}
idx_to_ctg = {}
ctg_to_idx = {}
fin = open(node_file, "rb")
for line in fin:
	# tokenize.
	tmp = line.replace("\n","").split("\t")
	
	# save info.
	idx_to_ctg[tmp[0]] = tmp[1]
	ctg_to_idx[tmp[1]] = tmp[0]
	
	# choose flip status.
	if random.random() > 0.5:
		flip_status[tmp[0]] = True
	else:
		flip_status[tmp[0]] = False
		
fin.close()

def randomize_edges(fs, efile, ofile):
	
	# read edges.
	fin = open(efile, "rb")
	lines = fin.readlines()
	fin.close()

	fout = open(ofile, "wb")
	for line in lines:
		# tokenize.
		tmp = line.replace("\n","").split("\t")
		ctg1 = tmp[0]
		ctg2 = tmp[1]

		# see if we flipped any of the contigs.
		if fs[ctg1] == True:
			if tmp[6] == "0":
				tmp[6] = "1"
			else:
				tmp[6] = "0"

		if fs[ctg2] == True:
			if tmp[7] == "0":
				tmp[7] = "1"
			else:
				tmp[7] = "0"
	 
		# print edges.
		fout.write('\t'.join(tmp) + "\n" )
	fout.close()

# randomize edges.
randomize_edges(flip_status, edge_file1, out2)
randomize_edges(flip_status, edge_file2, out3)
	
# update the agp.
nodes = []
fin = open(agp_file, "rb")
fout = open(out1, "wb")
for line in fin:
	# tokenize.
	if line[0] == "#": continue
	tmp = line.replace("\n","").split("\t")

	# check if this was flipped.
	id = ctg_to_idx[tmp[5]]
	
	if flip_status[id] == True:
		if tmp[8] == "+":
			tmp[8] = "-"
		else:
			tmp[8] = "+"
	
	# write results to agp.
	fout.write('\t'.join(tmp) + "\n")
	
fout.close()
fin.close()
